// Copyright 2016 The kingshard Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package backend

import (
	"context"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jackc/pgproto3/v2"

	"gitee.com/he3db/he3proxy/config"
	"gitee.com/he3db/he3proxy/core/golog"

	"gitee.com/he3db/he3proxy/core/errors"
	"gitee.com/he3db/he3proxy/mysql"
)

const (
	Up = iota
	Down
	ManualDown
	Unknown

	InitConnCount           = 16
	DefaultMaxConnNum       = 1024
	PingPeroid        int64 = 4
	MaxPoolNum              = 4
)

type DB struct {
	sync.RWMutex

	addr     string
	user     string
	password string
	db       string
	state    int32

	maxConnNum  int
	InitConnNum int
	idleConns   chan *Conn
	cacheConns  chan *Conn
	checkConn   *Conn
	lastPing    int64

	pushConnCount int64
	popConnCount  int64

	//pg
	cacheConnsMap map[string]chan *Conn
	currConnCount int64 //当前总链接数
	maxPoolNum    int   //支持的最大db链接池
}

func Open(addr string, user string, password string, dbName string, maxConnNum int) (*DB, error) {
	var err error
	db := new(DB)
	db.addr = addr
	db.user = user
	db.password = password
	db.db = dbName

	if 0 < maxConnNum {
		db.maxConnNum = maxConnNum
		if db.maxConnNum < 16 {
			db.InitConnNum = db.maxConnNum
		} else {
			db.InitConnNum = db.maxConnNum / 4
		}
	} else {
		db.maxConnNum = DefaultMaxConnNum
		db.InitConnNum = InitConnCount
	}
	//check connection
	db.checkConn, err = db.newConn("")
	if err != nil {
		db.Close()
		return nil, err
	}

	//因pg链接通道与db绑定，所以只有mysql才在启动时设置链接池，pg链接池将在客户端链接请求进来时初始化
	if config.Mysql == os.Getenv(config.DbTypeEnv) {
		db.idleConns = make(chan *Conn, db.maxConnNum)
		db.cacheConns = make(chan *Conn, db.maxConnNum)
		atomic.StoreInt32(&(db.state), Unknown)

		for i := 0; i < db.maxConnNum; i++ {
			if i < db.InitConnNum {
				conn, err := db.newConn("")

				if err != nil {
					db.Close()
					return nil, err
				}

				db.cacheConns <- conn
				atomic.AddInt64(&db.pushConnCount, 1)
			} else {
				conn := new(Conn)
				db.idleConns <- conn
				atomic.AddInt64(&db.pushConnCount, 1)
			}
		}
	}
	db.SetLastPing()

	return db, nil
}

func (db *DB) Addr() string {
	return db.addr
}

func (db *DB) State() string {
	var state string
	switch db.state {
	case Up:
		state = "up"
	case Down, ManualDown:
		state = "down"
	case Unknown:
		state = "unknow"
	}
	return state
}

func (db *DB) ConnCount() (int, int, int64, int64) {
	db.RLock()
	defer db.RUnlock()
	return len(db.idleConns), len(db.cacheConns), db.pushConnCount, db.popConnCount
}

func (db *DB) Close() error {
	db.Lock()
	idleChannel := db.idleConns
	cacheChannel := db.cacheConns
	db.cacheConns = nil
	db.idleConns = nil
	db.Unlock()
	if cacheChannel == nil || idleChannel == nil {
		return nil
	}

	close(cacheChannel)
	for conn := range cacheChannel {
		db.closeConn(conn)
	}
	close(idleChannel)

	return nil
}

func (db *DB) getConns() (chan *Conn, chan *Conn) {
	db.RLock()
	cacheConns := db.cacheConns
	idleConns := db.idleConns
	db.RUnlock()
	return cacheConns, idleConns
}

func (db *DB) getCacheConns() chan *Conn {
	db.RLock()
	conns := db.cacheConns
	db.RUnlock()
	return conns
}

func (db *DB) getIdleConns() chan *Conn {
	db.RLock()
	conns := db.idleConns
	db.RUnlock()
	return conns
}

func (db *DB) Ping() error {
	var err error
	if db.checkConn == nil {
		db.checkConn, err = db.newConn("")
		if err != nil {
			if db.checkConn != nil {
				db.checkConn.Close()
				db.checkConn = nil
			}
			return err
		}
	}
	err = db.checkConn.Ping()
	if err != nil {
		if db.checkConn != nil {
			db.checkConn.Close()
			db.checkConn = nil
		}
		return err
	}
	return nil
}

func (db *DB) newConn(user string) (*Conn, error) {
	co := new(Conn)

	if user == "" {
		user = db.user
	}

	if err := co.Connect(db.addr, user, db.password, db.db); err != nil {
		return nil, err
	}

	co.pushTimestamp = time.Now().Unix()

	return co, nil
}

func (db *DB) addIdleConn() {
	conn := new(Conn)
	select {
	case db.idleConns <- conn:
	default:
		break
	}
}

func (db *DB) closeConn(co *Conn) error {
	atomic.AddInt64(&db.pushConnCount, 1)
	if config.Mysql == os.Getenv(config.DbTypeEnv) {
		if co != nil {
			co.Close()
			conns := db.getIdleConns()
			if conns != nil {
				select {
				case conns <- co:
					return nil
				default:
					return nil
				}
			}
		} else {
			db.addIdleConn()
		}
	} else {
		if co != nil {
			// If the user link is not cached in the connection pool, it will be released directly.
			// If the link is cached, the current link will be released and a new link will be generated to join the pool.
			// The number of connections in the maintenance pool will remain unchanged
			conns := db.cacheConnsMap[db.user+db.db]
			if conns == nil || len(conns) == db.InitConnNum {
				atomic.AddInt64(&db.currConnCount, -1)
			} else {
				conn, err := db.newConn(db.user)
				if err != nil {
					return nil
				}
				conns <- conn
				atomic.AddInt64(&db.pushConnCount, 1)
			}
			co.Close()
		}
	}

	return nil
}

func (db *DB) closeConnNotAdd(co *Conn) error {
	if co != nil {
		co.Close()
		conns := db.getIdleConns()
		if conns != nil {
			select {
			case conns <- co:
				return nil
			default:
				return nil
			}
		}
	} else {
		db.addIdleConn()
	}
	return nil
}

func (db *DB) tryReuse(co *Conn) error {
	var err error
	//reuse Connection
	if co.IsInTransaction() {
		//we can not reuse a connection in transaction status
		err = co.Rollback()
		if err != nil {
			return err
		}
	}

	if !co.IsAutoCommit() {
		//we can not  reuse a connection not in autocomit
		_, err = co.exec("set autocommit = 1")
		if err != nil {
			return err
		}
	}

	//connection may be set names early
	//we must use default utf8
	if co.GetCharset() != mysql.DEFAULT_CHARSET {
		err = co.SetCharset(mysql.DEFAULT_CHARSET, mysql.DEFAULT_COLLATION_ID)
		if err != nil {
			return err
		}
	}

	return nil
}

func (db *DB) PopConn() (*Conn, error) {
	var co *Conn
	var err error

	cacheConns, idleConns := db.getConns()
	if cacheConns == nil || idleConns == nil {
		return nil, errors.ErrDatabaseClose
	}
	co = db.GetConnFromCache(cacheConns)
	if co == nil {
		co, err = db.GetConnFromIdle(cacheConns, idleConns)
		if err != nil {
			return nil, err
		}
	}

	err = db.tryReuse(co)
	if err != nil {
		db.closeConn(co)
		return nil, err
	}

	return co, nil
}

func (db *DB) PopConnPg(dbname string, dbuser string) (*Conn, error) {
	var co *Conn
	var err error

	db.db = dbname
	db.user = dbuser

	//判断是否关闭链接池 默认开启
	if os.Getenv(config.ConnPoolSwitch) == "false" {
		//返回一个链接
		co, err = db.newConn(dbuser)
		if err != nil {
			db.Close()
			return nil, err
		}
	} else {
		if db.cacheConnsMap == nil {
			db.cacheConnsMap = make(map[string]chan *Conn)
		}
		cacheConns := db.cacheConnsMap[dbuser+dbname]
		if cacheConns == nil {
			if db.IsExceedMaxConns() {
				err = errors.ErrConnIsFull
				return nil, err
			}
			//判断是否超过最大链接池数量
			if len(db.cacheConnsMap) >= db.maxPoolNum {
				err = errors.ErrMaxPoolIsFull
				return nil, err
			}
			//InitConnNum 根据配置文件配置而来；默认值16
			cacheConns = make(chan *Conn, db.InitConnNum)
			atomic.StoreInt32(&(db.state), Unknown)

			//循环建立链接时间较慢 暂且使用协程处理;
			//2022.06.08 测试发现并发下 协程创建 会导致链接池数量不稳定，因此暂时注销，采用同步方式，这样第一次链接查询就会变慢
			//go func() {
			//初始化线程池
			for i := 0; i < db.InitConnNum; i++ {
				conn, err := db.newConn(dbuser)

				if err != nil {
					db.Close()
				}

				cacheConns <- conn
				atomic.AddInt64(&db.pushConnCount, 1)
				atomic.AddInt64(&db.currConnCount, 1)
			}
			db.cacheConnsMap[dbuser+dbname] = cacheConns
			//}()

			//返回一个链接
			//co, err = db.newConn(dbuser)
			//if err != nil {
			//	db.Close()
			//	return nil, err
			//}
			//atomic.AddInt64(&db.pushConnCount, 1)
			//atomic.AddInt64(&db.currConnCount, 1)
		}
		//else {
		co = db.GetConnFromCache(cacheConns)
		if co == nil {
			golog.Debug("db", "PopConnPg", "conn is nil", 0)
			// 需要判断是否大于最大链接数
			if db.IsExceedMaxConns() {
				err = errors.ErrConnIsFull
				return nil, err
			}
			//返回一个链接
			co, err = db.newConn(dbuser)
			if err != nil {
				db.Close()
				return nil, err
			}
			atomic.AddInt64(&db.pushConnCount, 1)
			atomic.AddInt64(&db.currConnCount, 1)
		}
		//}
	}

	//TODO 事物处理
	/*
		err = db.tryReuse(co)
		if err != nil {
			db.closeConn(co)
			return nil, err
		}*/

	return co, nil
}

//TODO 如何判断是否超过最大链接数
func (db *DB) IsExceedMaxConns() bool {
	//strInt64 := strconv.FormatInt(db.currConnCount, 10)
	//curr ,_ := strconv.Atoi(strInt64)
	//if curr >= db.maxConnNum {
	//	return true
	//}else {
	//	return false
	//}
	return false
}

func (db *DB) GetConnFromCache(cacheConns chan *Conn) *Conn {
	var co *Conn
	var err error
	for 0 < len(cacheConns) {
		co = <-cacheConns
		atomic.AddInt64(&db.popConnCount, 1)
		if co != nil && PingPeroid < time.Now().Unix()-co.pushTimestamp {
			err = co.Ping()
			if err != nil {
				db.closeConn(co)
				co = nil
			}
		}
		if co != nil {
			break
		}
	}
	return co
}

func (db *DB) GetConnFromIdle(cacheConns, idleConns chan *Conn) (*Conn, error) {
	var co *Conn
	var err error
	select {
	case co = <-idleConns:
		atomic.AddInt64(&db.popConnCount, 1)
		co, err := db.newConn("")
		if err != nil {
			db.closeConn(co)
			return nil, err
		}
		err = co.Ping()
		if err != nil {
			db.closeConn(co)
			return nil, errors.ErrBadConn
		}
		return co, nil
	case co = <-cacheConns:
		atomic.AddInt64(&db.popConnCount, 1)
		if co == nil {
			return nil, errors.ErrConnIsNil
		}
		if co != nil && PingPeroid < time.Now().Unix()-co.pushTimestamp {
			err = co.Ping()
			if err != nil {
				db.closeConn(co)
				return nil, errors.ErrBadConn
			}
		}
	}
	return co, nil
}

func (db *DB) PushConnForExtendedProtocol(co *Conn, err error) {
	if co == nil {
		return
	}
	db.RLock()
	conns := db.cacheConnsMap[db.user+db.db]
	db.RUnlock()
	// 因后端链接未销毁 到时扩展协议时报错： prepared statement "lrupsc_1_0" already exists (SQLSTATE 42P05)
	// 如果成功创建了一个命名的预备语句对象，那么它将持续到当前会话结束， 除非被明确地删除 （暂未实现 需要采用close命令关闭），现在是直接关闭了链接
	// http://www.postgres.cn/docs/14/protocol-message-formats.html
	co.Close()
	if conns == nil || err != nil {
		return
	}
	co.pushTimestamp = time.Now().Unix()
	co, err = db.newConn("")
	if err != nil {
		return
	}
	select {
	case conns <- co:
		return
	default:
		co.Close()
		return
	}
}

func (db *DB) PushConn(co *Conn, err error) {
	atomic.AddInt64(&db.pushConnCount, 1)

	if config.Mysql == os.Getenv(config.DbTypeEnv) {
		if co == nil {
			db.addIdleConn()
			return
		}
		conns := db.getCacheConns()
		if conns == nil {
			co.Close()
			return
		}
		if err != nil {
			db.closeConnNotAdd(co)
			return
		}
		co.pushTimestamp = time.Now().Unix()
		select {
		case conns <- co:
			return
		default:
			db.closeConnNotAdd(co)
			return
		}
	} else {
		if co == nil {
			return
		}
		db.RLock()
		conns := db.cacheConnsMap[db.user+db.db]
		db.RUnlock()
		if conns == nil || err != nil || len(conns) == db.InitConnNum {
			co.Close()
			atomic.AddInt64(&db.currConnCount, -1)
			return
		}
		co.pushTimestamp = time.Now().Unix()
		select {
		case conns <- co:
			return
		default:
			co.Close()
			atomic.AddInt64(&db.currConnCount, -1)
			return
		}
	}
}

type BackendConn struct {
	*Conn
	db              *DB
	ParName         []string //parse 通道name
	IsInTransaction bool     // 是否在事务中
}

func (p *BackendConn) Close() {
	if p != nil && p.Conn != nil {
		// 因事务中的链接存在较多数据 如临时表、listen等， 随seesion端开而消失，因此直接该后端链接释放
		if p.Conn.pkgErr != nil || p.IsInTransaction {
			p.db.closeConn(p.Conn)
		} else {
			var data []byte
			if p.ParName != nil {
				for _, name := range p.ParName {
					// 'S'关闭一个准备的语句；或者'P'关闭一个入口。
					// 结合之前pgx链接报错 prepared statement "lrupsc_1_0" already exists (SQLSTATE 42P05)
					// 因此这里应该使用S  进行关闭
					c := pgproto3.Close{'S', name}
					data = c.Encode(data)
					//c1 := pgproto3.Close{'P', name}
					//data = c1.Encode(data)
				}
				f := pgproto3.Flush{}
				data = f.Encode(data)
				p.ParName = nil
			}
			//reset := pgproto3.Query{"RESET ALL"}
			//data = reset.Encode(data)
			//err := p.Conn.WritePgPacket(data)
			//p.Conn.ReadPgPacket()
			_, err := p.Conn.ConnPg.Exec(context.Background(), "RESET ALL")
			if err != nil {
				golog.Warn("db", "Close", "write close parse msg err: "+err.Error(), '0')
			}
			p.db.PushConn(p.Conn, nil)
		}
		p.Conn = nil
	}
}

func (p *BackendConn) CloseForExtendedProtocol() {
	if p != nil && p.Conn != nil {
		if p.Conn.pkgErr != nil {
			p.db.closeConn(p.Conn)
		} else {
			go p.db.PushConnForExtendedProtocol(p.Conn, nil)
		}
		p.Conn = nil
	}
}

func (db *DB) GetConn() (*BackendConn, error) {
	c, err := db.PopConn()
	if err != nil {
		return nil, err
	}
	return &BackendConn{c, db, nil, false}, nil
}

func (db *DB) GetConnPg(dbname string, dbuser string) (*BackendConn, error) {
	c, err := db.PopConnPg(dbname, dbuser)
	if err != nil {
		return nil, err
	}
	return &BackendConn{c, db, nil, false}, nil
}

func (db *DB) SetLastPing() {
	db.lastPing = time.Now().Unix()
}

func (db *DB) GetLastPing() int64 {
	return db.lastPing
}
